require(tidyverse)
require(ggrepel)
rm(list = ls())
gc()

################################################################################
##
## Date:    2024-12-20
## Author:  james.h.bisbee@vanderbilt.edu
## Purpose: This script generates SI Figure 1.
## Inputs:  ./data/VIMP_lasso/LASSO_2024_outcome-SOLSAT_period-years_temp-chg12.RData
##          ./data/VIMP_lasso/LASSO_2024_outcome-ECON_period-years_temp-chg12.RData
##            - LASSO results generated by NFR_lasso_prep.R
## Outputs: ./figures/figS1.pdf
##
################################################################################

# Compute details
print(paste0('Compute environment from ',Sys.Date(),' run by Bisbee'))
if(Sys.info()['sysname'] == 'Windows') {
  ram_size = system("wmic MemoryChip get Capacity", intern = TRUE)[-1]
  model_name = system("wmic cpu get name", intern = TRUE)[2] # nocov
  vendor_id = system("wmic cpu get manufacturer", intern = TRUE)[2] # nocov
  
  print(list(ram = stringr::str_squish(ram_size)[1],
             vendor_id = stringr::str_squish(vendor_id),
             model_name = stringr::str_squish(model_name),
             no_of_cores = parallel::detectCores()))
} else if(Sys.info()['sysname'] == 'Linuxs') {
  splitted <- strsplit(system("ps -C rsession -o %cpu,%mem,pid,cmd", intern = TRUE), " ")
  df <- do.call(rbind, lapply(splitted[-1], 
                              function(x) data.frame(
                                cpu = as.numeric(x[2]),
                                mem = as.numeric(x[4]),
                                pid = as.numeric(x[5]),
                                cmd = paste(x[-c(1:5)], collapse = " "))))
  df
} else {
  cat("If not on Linux or Windows, you'll have to figure out your own solution to seeing the compute environment.")
}

sessionInfo()

load('./data/VIMP_lasso/LASSO_2024_outcome-SOLSAT_period-years_temp-chg12.RData')

toplot <- lassoRes %>%
  left_join(lambdaLookup %>%
              mutate(normInd = as.integer(gsub('s','',normInd)))) %>%
  group_by(vars,period,normInd,outcome) %>%
  summarise(lambda = mean(lambda),
            norm = mean(norm),
            dev.ratio = mean(dev.ratio),
            m = median(coef),
            lb = quantile(coef,.025),
            ub = quantile(coef,.975),.groups = 'drop')

toplot <- toplot %>%
  left_join(coefRes %>%
              group_by(vars,outcome,period) %>%
              summarise(propIncl=n()/100,.groups = 'drop') %>%
              mutate(vars = gsub('(AGECAT|EDUC|GENDER|INC|MARST|PARTY|RACE)','\\1_',vars)),
            by = c('vars','period','outcome'))

load('./data/VIMP_lasso/LASSO_2024_outcome-ECON_period-years_temp-chg12.RData')

tmp <- lassoRes %>%
  left_join(lambdaLookup %>%
              mutate(normInd = as.integer(gsub('s','',normInd)))) %>%
  group_by(vars,period,normInd,outcome) %>%
  summarise(lambda = mean(lambda),
            norm = mean(norm),
            dev.ratio = mean(dev.ratio),
            m = median(coef),
            lb = quantile(coef,.025),
            ub = quantile(coef,.975),.groups = 'drop')

tmp <- tmp %>%
  left_join(coefRes %>%
              group_by(vars,outcome,period) %>%
              summarise(propIncl=n()/100,.groups = 'drop') %>%
              mutate(vars = gsub('(AGECAT|EDUC|GENDER|INC|MARST|PARTY|RACE)','\\1_',vars)),
            by = c('vars','period','outcome'))

toplot <- toplot %>%
  bind_rows(tmp)

lookup <- toplot %>%
  count(vars) %>%
  mutate(cat = gsub('_.*','',vars)) %>%
  mutate(type = ifelse(cat %in% c('AGECAT','EDUC','GENDER','INC','MARST','PARTY','RACE'),'Indiv','Geo')) %>%
  mutate(levels = gsub('.*_?','',vars))

lookup <- lookup %>%
  mutate(levels = str_to_title(gsub('Living wi','',gsub('^ur$','Unemp Rate',gsub('lfpr','LF Participation Rate',gsub('_',' ',gsub('(\\d+)_(\\d+)','\\1-\\2',gsub('^[A-Z]+_','',gsub('ARRST_','',gsub('imputed_|_imputed','',vars)))))))))) %>%
  mutate(catLabs = ifelse(cat == 'AGECAT','AGE',
                          ifelse(cat == 'ECON','LBR',cat))) %>%
  mutate(levels= paste0(catLabs,': ',levels))

toplot2 <- toplot %>%
  filter(period %in% c('2016-01-01'),!grepl('_emplvl|livexp|REF',vars)) %>%
  group_by(vars,period,outcome) %>%
  mutate(minInd = min(normInd)) %>%
  left_join(lookup) %>%
  ungroup() %>%
  mutate(Type = paste0(type,'-',catLabs)) %>%
  mutate(outcome = ifelse(outcome == 'ECON','Views of the economy','Satisfaction with standard of living')) %>%
  mutate(Type = ifelse(Type %in% c('Indiv-AGE','Indiv-GENDER','Indiv-RACE'),'Indiv: ASR',
                       ifelse(Type %in% c('Indiv-EDUC','Indiv-INC','Indiv-MARST'),'Indiv: Educ, Inc & Marital',
                              ifelse(Type %in% c('Indiv-PARTY'),'Indiv: Partisanship',str_to_title(gsub('-',': ',Type))))))


xlabs <- toplot2 %>%
  select(normInd,lambda,norm) %>%
  distinct() %>%
  group_by(normInd) %>%
  summarise_all(mean) %>%
  arrange(normInd)


toplot3 <- toplot2 %>%
  drop_na(propIncl) %>%
  mutate(outcome = ifelse(outcome == 'Satisfaction with standard of living',
                          'Are you satisfied or dissatisfied with\nyour standard of living?',
                          'How would you rate economic conditions\nin this country today?'))

pdf('./figures/figS1.pdf',width = 7,height = 5)
toplot3 %>%
  ggplot(aes(x = normInd,y = m,group = vars,color = Type,alpha = propIncl == 1)) + 
  geom_line() + 
  geom_text_repel(data = toplot3 %>% 
                    group_by(outcome) %>%
                    filter(normInd == max(normInd),propIncl ==1 ) %>%
                    drop_na(propIncl),aes(label = levels),hjust = 0,direction = 'y',
                  size = 2.5) +
  theme_bw() + 
  scale_alpha_manual(guide = 'none',values = c(.05,1)) + 
  scale_x_continuous(expand = c(.05,0,.5,0),breaks = c(1,10,20,29),
                     labels = round(xlabs %>% filter(normInd %in% c(1,10,20,28)) %>% .$lambda,4)) +
  labs(
    x = 'Lambda Penalty',
    y = 'Coefficient',
    title = 'LASSO Results') + 
  facet_grid(~outcome) + 
  geom_hline(yintercept = 0,linetype = 'dashed') + 
  theme(legend.position = 'bottom')
dev.off()

# EOF